from datetime import datetime
import uuid
import json
from app import db
from app.models.user import User  # 导入User模型

class Scan(db.Model):
    """扫描任务表"""
    __tablename__ = 'scans'
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    user_id = db.Column(db.String(36), db.ForeignKey('users.id'), nullable=True)
    target = db.Column(db.String(255), nullable=False)
    status = db.Column(db.String(20), default='pending')  # pending, running, completed, failed
    scan_depth = db.Column(db.Integer, default=2)
    scan_types = db.Column(db.Text, default='[]')
    scheduled = db.Column(db.Boolean, default=False)
    schedule_id = db.Column(db.String(36), nullable=True)
    tags = db.Column(db.Text, default='[]')  # JSON格式存储标签
    is_public = db.Column(db.Boolean, default=False)  # 是否公开共享
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    finished_at = db.Column(db.DateTime, nullable=True)
    
    # 关系：一个扫描任务可以有多个漏洞
    vulnerabilities = db.relationship('Vulnerability', backref='scan', lazy=True)

class Vulnerability(db.Model):
    """漏洞信息表"""
    __tablename__ = 'vulnerabilities'
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    scan_id = db.Column(db.String(36), db.ForeignKey('scans.id'), nullable=False)
    vuln_id = db.Column(db.String(50), nullable=False)  # 漏洞唯一标识符，如VULN-2025-0001
    vuln_type = db.Column(db.String(50), nullable=False)
    path = db.Column(db.String(255), nullable=False)
    payload = db.Column(db.Text, nullable=True)
    evidence = db.Column(db.Text, nullable=True)
    raw_score = db.Column(db.Float, nullable=True)  # 原始分数
    risk_score = db.Column(db.Float, nullable=True)  # 风险评分(0-1)
    severity = db.Column(db.String(20), nullable=True)  # low, medium, high
    suggestion = db.Column(db.Text, nullable=True)  # 修复建议
    status = db.Column(db.String(20), default='unhandled')  # unhandled, confirmed, in_progress, fixed, verified
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    def to_dict(self):
        """将漏洞对象转换为字典格式"""
        return {
            'id': self.id,
            'scan_id': self.scan_id,
            'vuln_id': self.vuln_id,
            'vuln_type': self.vuln_type,
            'path': self.path,
            'payload': self.payload,
            'evidence': self.evidence,
            'raw_score': self.raw_score,
            'risk_score': self.risk_score,
            'severity': self.severity,
            'status': self.status,
            'suggestion': self.suggestion,
            'created_at': self.created_at.isoformat(),
            'updated_at': self.updated_at.isoformat()
        }

# 多对多关联表：漏洞与修复建议
vulnerability_suggestions = db.Table('vulnerability_suggestions',
    db.Column('vulnerability_id', db.String(36), db.ForeignKey('vulnerabilities.id'), primary_key=True),
    db.Column('suggestion_id', db.String(36), db.ForeignKey('repair_suggestions.id'), primary_key=True)
)

class RepairSuggestion(db.Model):
    """修复建议表"""
    __tablename__ = 'repair_suggestions'
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    title = db.Column(db.String(255), nullable=False)
    description = db.Column(db.Text, nullable=True)
    fix_steps = db.Column(db.Text, default='[]')  # JSON格式存储修复步骤
    code_example = db.Column(db.Text, nullable=True)
    best_practices = db.Column(db.Text, default='[]')  # JSON格式存储最佳实践
    recommended_tools = db.Column(db.Text, default='[]')  # JSON格式存储推荐工具
    source = db.Column(db.String(50), default='system')  # system, ai, manual
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    last_updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    effectiveness_rating = db.Column(db.Float, nullable=True)
    
    # 关系
    vulnerabilities = db.relationship('Vulnerability', secondary=vulnerability_suggestions, 
                                     backref=db.backref('repair_suggestions', lazy='dynamic'))
    
    def to_dict(self):
        """转换为字典格式"""
        return {
            'id': self.id,
            'title': self.title,
            'description': self.description,
            'fix_steps': json.loads(self.fix_steps) if self.fix_steps else [],
            'code_example': self.code_example,
            'best_practices': json.loads(self.best_practices) if self.best_practices else [],
            'recommended_tools': json.loads(self.recommended_tools) if self.recommended_tools else [],
            'source': self.source,
            'created_at': self.created_at.isoformat(),
            'last_updated_at': self.last_updated_at.isoformat(),
            'effectiveness_rating': self.effectiveness_rating
        }

class RiskAssessment(db.Model):
    """风险评估表"""
    __tablename__ = 'risk_assessments'
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    scan_id = db.Column(db.String(36), db.ForeignKey('scans.id'), nullable=True)
    vulnerability_id = db.Column(db.String(36), db.ForeignKey('vulnerabilities.id'), nullable=True)
    
    # 风险评分
    risk_score = db.Column(db.Float, nullable=False)  # 风险总分
    risk_level = db.Column(db.String(50), nullable=False)  # high, medium, low
    
    # 评估因素
    exploitation_probability = db.Column(db.Float, default=0.5)
    potential_impact = db.Column(db.Float, default=0.5)
    data_sensitivity = db.Column(db.Float, default=0.5)
    business_criticality = db.Column(db.Float, default=0.5)
    
    # 特征向量
    feature_vector = db.Column(db.Text, default='{}')  # JSON格式存储特征向量
    model_version = db.Column(db.String(50), nullable=True)
    prediction_confidence = db.Column(db.Float, default=0.9)
    
    # 时间信息
    assessed_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    def to_dict(self):
        """转换为字典格式"""
        return {
            'id': self.id,
            'scan_id': self.scan_id,
            'vulnerability_id': self.vulnerability_id,
            'risk_score': self.risk_score,
            'risk_level': self.risk_level,
            'exploitation_probability': self.exploitation_probability,
            'potential_impact': self.potential_impact,
            'data_sensitivity': self.data_sensitivity,
            'business_criticality': self.business_criticality,
            'feature_vector': json.loads(self.feature_vector) if self.feature_vector else {},
            'model_version': self.model_version,
            'prediction_confidence': self.prediction_confidence,
            'assessed_at': self.assessed_at.isoformat()
        }

class AttackPath(db.Model):
    """攻击路径表"""
    __tablename__ = 'attack_paths'
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    scan_id = db.Column(db.String(36), db.ForeignKey('scans.id'), nullable=False)
    
    # 路径信息
    path_name = db.Column(db.String(255), nullable=False)
    path_score = db.Column(db.Float, nullable=False)
    path_level = db.Column(db.String(50), nullable=False)
    
    # 路径步骤
    path_steps = db.Column(db.Text, default='[]')  # JSON格式存储攻击步骤
    path_length = db.Column(db.Integer, default=0)
    
    # 攻击目标
    target_host = db.Column(db.String(45), nullable=True)
    target_service = db.Column(db.String(100), nullable=True)
    target_asset = db.Column(db.String(255), nullable=True)
    
    # 技术信息
    attack_techniques = db.Column(db.Text, default='[]')  # JSON格式存储攻击技术
    required_privileges = db.Column(db.String(100), nullable=True)
    difficulty = db.Column(db.String(50), default='medium')
    
    # 时间信息
    discovered_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    def to_dict(self):
        """转换为字典格式"""
        return {
            'id': self.id,
            'scan_id': self.scan_id,
            'path_name': self.path_name,
            'path_score': self.path_score,
            'path_level': self.path_level,
            'path_steps': json.loads(self.path_steps) if self.path_steps else [],
            'path_length': self.path_length,
            'target_host': self.target_host,
            'target_service': self.target_service,
            'target_asset': self.target_asset,
            'attack_techniques': json.loads(self.attack_techniques) if self.attack_techniques else [],
            'required_privileges': self.required_privileges,
            'difficulty': self.difficulty,
            'discovered_at': self.discovered_at.isoformat()
        }

class TaskStatus(db.Model):
    """任务状态跟踪表（可选，用于Celery任务跟踪）"""
    __tablename__ = 'task_status'
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    task_id = db.Column(db.String(50), unique=True, nullable=False)
    scan_id = db.Column(db.String(36), db.ForeignKey('scans.id'), nullable=False)
    status = db.Column(db.String(20), default='pending')
    progress = db.Column(db.Integer, default=0)
    result = db.Column(db.Text, nullable=True)
    error = db.Column(db.Text, nullable=True)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

class ScanResult(db.Model):
    """扫描结果详情表"""
    __tablename__ = 'scan_results'
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    scan_id = db.Column(db.String(36), db.ForeignKey('scans.id'), nullable=False)
    
    # 扫描元数据
    scan_date = db.Column(db.DateTime, default=datetime.utcnow)
    scan_duration = db.Column(db.Float, nullable=True)
    scan_status = db.Column(db.String(50), default='completed')
    
    # 发现信息
    discovered_urls = db.Column(db.Text, default='[]')  # JSON格式
    discovered_forms = db.Column(db.Text, default='[]')  # JSON格式
    discovered_endpoints = db.Column(db.Text, default='[]')  # JSON格式
    
    # 安全检查结果
    security_headers = db.Column(db.Text, default='{}')  # JSON格式
    ssl_tls_info = db.Column(db.Text, default='{}')  # JSON格式
    cookie_analysis = db.Column(db.Text, default='{}')  # JSON格式
    
    # 其他信息
    notes = db.Column(db.Text, nullable=True)
    scan_metadata = db.Column(db.Text, default='{}')  # JSON格式，避免使用保留关键字
    
    def to_dict(self):
        """转换为字典格式"""
        return {
            'id': self.id,
            'scan_id': self.scan_id,
            'scan_date': self.scan_date.isoformat(),
            'scan_duration': self.scan_duration,
            'scan_status': self.scan_status,
            'discovered_urls': json.loads(self.discovered_urls) if self.discovered_urls else [],
            'discovered_forms': json.loads(self.discovered_forms) if self.discovered_forms else [],
            'discovered_endpoints': json.loads(self.discovered_endpoints) if self.discovered_endpoints else [],
            'security_headers': json.loads(self.security_headers) if self.security_headers else {},
            'ssl_tls_info': json.loads(self.ssl_tls_info) if self.ssl_tls_info else {},
            'cookie_analysis': json.loads(self.cookie_analysis) if self.cookie_analysis else {},
            'notes': self.notes,
            'metadata': json.loads(self.scan_metadata) if self.scan_metadata else {}
        }

class TargetRange(db.Model):
    """真实授权靶场配置表"""
    __tablename__ = 'target_ranges'
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    user_id = db.Column(db.String(36), db.ForeignKey('users.id'), nullable=True)
    name = db.Column(db.String(100), nullable=False)
    base_url = db.Column(db.String(255), nullable=False)
    api_key = db.Column(db.String(255), nullable=True)  # 敏感信息加密存储
    auth_type = db.Column(db.String(30), default='api_key')  # api_key, oauth2, basic
    username = db.Column(db.String(100), nullable=True)
    password = db.Column(db.String(255), nullable=True)  # 敏感信息加密存储
    description = db.Column(db.Text, nullable=True)
    tags = db.Column(db.Text, default='[]')  # JSON格式存储标签
    scan_frequency = db.Column(db.String(50), default='manual')  # manual, daily, weekly, monthly
    is_active = db.Column(db.Boolean, default=True)
    last_connected = db.Column(db.DateTime, nullable=True)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    def to_dict(self):
        """将靶场配置对象转换为字典格式（不包含敏感信息）"""
        import json
        return {
            'id': self.id,
            'user_id': self.user_id,
            'name': self.name,
            'base_url': self.base_url,
            'auth_type': self.auth_type,
            'description': self.description,
            'tags': json.loads(self.tags) if self.tags else [],
            'scan_frequency': self.scan_frequency,
            'is_active': self.is_active,
            'last_connected': self.last_connected.isoformat() if self.last_connected else None,
            'created_at': self.created_at.isoformat(),
            'updated_at': self.updated_at.isoformat()
        }
        
    def to_dict_with_credentials(self):
        """将靶场配置对象转换为字典格式（包含敏感信息，仅在内部使用）"""
        result = self.to_dict()
        result.update({
            'api_key': self.api_key,
            'username': self.username,
            'password': self.password
        })
        return result

class Schedule(db.Model):
    """扫描任务调度表"""
    __tablename__ = 'schedules'
    
    id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    user_id = db.Column(db.String(36), db.ForeignKey('users.id'), nullable=True)
    name = db.Column(db.String(100), nullable=False)
    targets = db.Column(db.Text, nullable=False)  # JSON格式存储目标列表
    scan_depth = db.Column(db.Integer, default=2)
    scan_types = db.Column(db.Text, nullable=False)  # JSON格式存储扫描类型列表
    cron_expression = db.Column(db.String(100), nullable=False)  # 调度表达式，如"0 0 * * *"表示每天午夜
    notification_email = db.Column(db.String(120), nullable=True)  # 报告发送邮箱
    report_format = db.Column(db.String(20), default='pdf')  # 报告格式
    next_run_at = db.Column(db.DateTime, nullable=True)
    last_run_at = db.Column(db.DateTime, nullable=True)
    is_active = db.Column(db.Boolean, default=True)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    def to_dict(self):
        """将调度任务对象转换为字典格式"""
        import json
        
        return {
            'id': self.id,
            'name': self.name,
            'targets': json.loads(self.targets),
            'scan_depth': self.scan_depth,
            'scan_types': json.loads(self.scan_types),
            'cron_expression': self.cron_expression,
            'next_run_at': self.next_run_at.isoformat() if self.next_run_at else None,
            'last_run_at': self.last_run_at.isoformat() if self.last_run_at else None,
            'is_active': self.is_active,
            'created_at': self.created_at.isoformat(),
            'updated_at': self.updated_at.isoformat()
        }